# 示例：create_react_agent_system_prompt.py (改造版)

"""
REACT + 系统提示词：
- REACT 代理：支持推理 + 调用工具。
- 系统提示词：通过在 messages 中加入 ("system", "提示内容") 实现，
  让代理按照指定规则或语言回复。
"""

import os
from dotenv import load_dotenv
from typing import Literal
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

# 自动加载 .env 文件中的 API key 等配置
load_dotenv()

# ----------------- 初始化 LLM -----------------
llm = ChatOpenAI(
    api_key=os.getenv("MODELSCOPE_API_KEY"),
    openai_api_base=os.getenv("MODELSCOPE_API_BASE"),
    model="Qwen/Qwen3-32B",
    streaming=True,
    temperature=0.7
)

# ----------------- 定义工具 -----------------
@tool
def get_weather(city: Literal["nyc", "sf"]):
    """根据城市返回预定义天气信息"""
    if city == "nyc":
        return "It might be cloudy in nyc"
    elif city == "sf":
        return "It's always sunny in sf"
    else:
        raise AssertionError("Unknown city")

tools = [get_weather]

# ----------------- 创建 REACT 代理 -----------------
from langgraph.prebuilt import create_react_agent

# 注意：不再使用 state_modifier
graph = create_react_agent(llm, tools=tools)

# ----------------- 定义打印函数 -----------------
def print_stream(stream):
    """
    打印流式输出
    - 支持 message 是 tuple 或对象
    """
    for s in stream:
        message = s["messages"][-1]
        if isinstance(message, tuple):
            print(message)
        else:
            message.pretty_print()

# ----------------- 执行示例 -----------------
inputs = {
    "messages": [
        ("system", "Respond in Chinese"),  # 系统提示词，控制输出语言
        ("user", "What's the weather in NYC?")
    ]
}

print_stream(graph.stream(inputs, stream_mode="values"))
